#include <sys/time.h>
#include "conv.h"
#include "im2col.h"
#include "utils.h"

int main()
{
    // Test parameters
    constexpr unsigned N_ICH = 4;
    constexpr unsigned N_OCH = 8;
    constexpr unsigned N_IH = 5;
    constexpr unsigned N_IW = 5;
    constexpr unsigned K = 3;
    constexpr unsigned P = 1;
    constexpr unsigned S = 1;
    constexpr unsigned N_OH = (N_IH + 2 * P - K) / S + 1;
    constexpr unsigned N_OW = (N_IW + 2 * P - K) / S + 1;
    constexpr unsigned VEC_LEN = N_OH * N_OW;
    constexpr unsigned P_ICH = 2;
    constexpr unsigned P_OCH = 4;
    constexpr unsigned BIT_ACTV = 8; 
    constexpr unsigned BIT_WGHT = 8; 
    constexpr unsigned BIT_CONV = 32;
    constexpr unsigned FOLD_I = N_ICH / P_ICH;
    constexpr unsigned FOLD_O = N_OCH / P_OCH;
    static_assert(N_ICH % P_ICH == 0, "N_ICH must be divisible by P_ICH");
    static_assert(N_OCH % P_OCH == 0, "N_OCH must be divisible by P_OCH");

    std::array<unsigned, 2> ifm_shape = {N_IH * N_IW, N_ICH};
    std::array<unsigned, 2> ofm_shape = {N_OH * N_OW, N_OCH};

    ap_int<BIT_WGHT> weight[N_OCH][K*K][N_ICH];
    for(unsigned oc = 0; oc < N_OCH; ++oc) {
        for(unsigned ic = 0; ic < N_ICH; ++ic) {
            for(unsigned k = 0; k < K*K; ++k) {
                ap_int<BIT_WGHT> val = oc*K*K*N_ICH + ic*K*K + k; // Small values for easy debugging
                weight[oc][k][ic] = val;
            }
        }
    }

    // Weight array
    static ap_uint<P_OCH * P_ICH * BIT_WGHT> weight_reshape[FOLD_O][FOLD_I][K*K];

    // Initialize weights with test data
    for (unsigned fo = 0; fo < FOLD_O; ++fo) {
        for (unsigned fi = 0; fi < FOLD_I; ++fi) {
            for (unsigned k = 0; k < K*K; ++k) {
                ap_uint<P_OCH * P_ICH * BIT_WGHT> w = 0;
                for (unsigned o = 0; o < P_OCH; ++o) {
                    for (unsigned i = 0; i < P_ICH; ++i) {
                        // Simple test pattern: weight = (fo + fi + k + o + i) % 256
                        unsigned oc = fo * P_OCH + o;
                        unsigned ic = fi * P_ICH + i;
                        ap_uint<BIT_WGHT> val = weight[oc][k][ic];
                        w(SLICE(BIT_WGHT, P_ICH * o + i)) = val;
                    }
                }
                weight_reshape[fo][fi][k] = w;
            }
        }
    }

    // Input stream
    data_stream<BIT_ACTV> s_input("s_input");
    data_stream<BIT_ACTV> s_input_golden("s_input_golden");

    gen_ifm<2>(s_input, ifm_shape);
    gen_ifm<2>(s_input_golden, ifm_shape);

    timeval start, end;
    gettimeofday(&start, NULL);

    data_stream<BIT_ACTV*P_ICH> s_input_e("s_input");
    expandWidth<BIT_ACTV, P_ICH * BIT_ACTV, N_IH * N_IW * N_ICH>(s_input, s_input_e);

    data_stream<BIT_ACTV*P_ICH> s_input_im2col("s_input_im2col");
    im2col<1, N_IH, N_IW, N_ICH, K, S, P, BIT_ACTV, P_ICH, 0>(s_input_e, s_input_im2col);

    data_stream<BIT_CONV*P_OCH> s_output_e("s_output_e");
    conv<P_ICH, P_OCH, N_ICH, N_OCH, K, BIT_ACTV, BIT_WGHT, BIT_CONV, N_OH*N_OW>(
        s_input_im2col, s_output_e, weight_reshape
    );
    data_stream<BIT_CONV> s_output("s_output");
    reduceWidth<BIT_CONV * P_OCH, BIT_CONV, N_OH * N_OW * FOLD_O>(s_output_e, s_output);
    gettimeofday(&end, NULL);

    data_stream<BIT_CONV> s_output_golden("s_output_golden");
    conv_golden<N_ICH, N_OCH, N_IH, N_IW, K, P, S, BIT_ACTV, BIT_WGHT, BIT_CONV>(
        s_input_golden, s_output_golden, weight
    );
    auto error_cnt = check_afm<BIT_CONV, 2>(s_output_golden, s_output, ofm_shape);
    if(error_cnt > 0) {
        std::cout << "Test Failed! " << error_cnt << " errors" << std::endl;
        return 1;
    } else {
        std::cout << "Test Passed! " << std::endl;
        return 0;
    }
}